library(doParallel)
library(foreach)
library(doFuture)

RNGkind("L'Ecuyer-CMRG")
set.seed(1234)


congs <- 101:113

# Calculate the number of cores
no_cores <-  min(detectCores() - 1, 10)


load("all_data.Rda")
load("legDataTmp.Rda")
load("speeches.Rda")
load("y.Rda")
load(paste0(getwd(), "/EMestimates.Rda"))


# Calculate the number of cores
no_cores <- min(detectCores() - 1, 10)

# define matrix inverse function
matrix_inverse <- function(X) chol2inv(chol(X))


##Set things up
N <- nrow(y)
J <- ncol(y)
C <- length(congs)

##Initial values using Heckman-Snyder

y <- as.matrix(y)
lower <- dataset$legis.data$icpsrLegis[grep("PELOSI",rownames(dataset$votes))]
upper <- dataset$legis.data$icpsrLegis[grep("SENSENBR",rownames(dataset$votes))]
lower <- which(rownames(y)==lower)
upper <- which(rownames(y)==upper)

#we need to make a vector of time periods
times <- as.numeric(gsub("-.*","",colnames(y)))-min(as.numeric(gsub("-.*","",colnames(y))))+1
TimeMat <- matrix(0,J,C)
for (j in 1:nrow(TimeMat)) TimeMat[j,times[j]] <- 1

#hyperprior parameters
b0 <- rep(0,2)
B0 <- 25*diag(2)
t0 <- 0
T0 <- 1
d0 <- rep(0,length(unique(all_data$cong)))
D0 <- rep(3,length(unique(all_data$cong)))
dd0 <- matrix(c(t0,d0))
DD0 <- diag(c(T0,D0))


# #definte the speech-time matrix
legTimeMatrix <- vector("list", N)
for (i in 1:nrow(y)) legTimeMatrix[[i]] <- matrix(rep(speeches[i,],each=C),J,C,byrow=TRUE)*TimeMat

#define the update functions


ab.update_chol <- function(j,Xstar,B0,Yhat){
  matrix_inverse(crossprod(Xstar) + matrix_inverse(B0))%*%(crossprod(Xstar,Yhat[,j]))
}


xd.update_chol <- function(i,be.update,Yhat,legTimeMatrix,DD0){
  Xstar <- cbind(be.update,legTimeMatrix[[i]])
  matrix_inverse(crossprod(Xstar) + matrix_inverse(DD0))%*%(crossprod(Xstar,Yhat[i,]))
}

##gonna start the bootstrap
mu_true <- -matrix(rep(results$al.current,N),N,J,byrow=TRUE) +
  outer(results$x.current,results$be.current) +
  as.matrix(results$de.current[,times]*speeches)

pr_true <- pnorm(mu_true)

##do the bootstraps

forFun <- list("pr_true"=pr_true, "results"=results, "N"=N, "J"=J, "C"=C,
               "times"=times, "legTimeMatrix"=legTimeMatrix, "speeches"=speeches,
               "DD0"=DD0, "B0"=B0, "upper"=upper, "lower"=lower)

all_draws <- NULL
for (n in 1:100) all_draws[[n]] <- matrix(runif(forFun$N*forFun$J),forFun$N,forFun$J)


boot_maker <- function(i) {
  
  #set.seed(1234)
  
  #generate bootstrap sample
  u_draws <- all_draws[[i]]
  y_imputed <- matrix(NA,forFun$N,forFun$J)
  y_imputed[forFun$pr_true>u_draws] <- 1
  y_imputed[forFun$pr_true<=u_draws] <- 0
  
  x.current <- forFun$results$x.current
  al.current <- forFun$results$al.current
  be.current <- forFun$results$be.current
  de.current <- forFun$results$de.current
  
  
  tol <- 0
  count <- 0
  deviance <- 100
  llik <- NULL
  
  ##create a bunch of matrix holders
  mu <- matrix(NA, forFun$N, forFun$J)
  dmu <- matrix(NA, forFun$N, forFun$J)
  pmu <- matrix(NA, forFun$N, forFun$J)
  
  while (tol<1-1e-3){
    
    iter.start <- proc.time()[3]
    
    ##create the NxJ matrix of legislator loyalty times speeches
    d_speech_mat <- as.matrix(de.current[,forFun$times]*forFun$speeches)
    
    ##calculate the matrix of means
    mu <- -matrix(rep(al.current,N),N,J,byrow=TRUE) +
      outer(x.current,be.current) +
      d_speech_mat
    
    #a minor correction for outliers
    mu[mu < -8] <- -8
    mu[mu > 8] <- 8
    
    #set the value of ystar for this iteration
    dmu <- dnorm(-mu)
    pmu <- pnorm(-mu)
    ystar <- (mu+dmu/(1-pmu))*(y_imputed==1) + (mu-dmu/(pmu))*(y_imputed==0)
    ystar[is.na(y_imputed)] <- mu[is.na(y_imputed)]
    
    
    ##estimate roll call parameters
    tmpRC <- sapply(1:forFun$J, ab.update_chol, Xstar=cbind(-1,x.current), B0=forFun$B0, Yhat=ystar-d_speech_mat)
    
    al.update <- tmpRC[1,]
    be.update <- tmpRC[2,]
    
    
    # Run across legislators
    tmpLegis <- sapply(1:forFun$N, xd.update_chol, be.update=tmpRC[2,], 
                       Yhat=ystar + matrix(rep(al.current,forFun$N),forFun$N,forFun$J,byrow=TRUE),
                       legTimeMatrix=forFun$legTimeMatrix,DD0=forFun$DD0)
    
    x.update <- tmpLegis[1,]
    de.update <- t(tmpLegis[2:nrow(tmpLegis),])
    
    
    ##make sure orientation of x-axis is correct
    if (x.update[forFun$lower]>x.update[forFun$upper]){
      x.update <- -x.update
      be.update <- -be.update
    }
    
    
    ##calculate correlation-based model tolerance
    if (count>0) tol <- min(c(cor(x.update,x.current), cor(al.update,al.current),
                              cor(be.update,be.current), cor(c(de.update),c(de.current))))
    
    
    ##set new values as those just estimated
    x.current <- x.update
    al.current <- al.update
    be.current <- be.update
    de.current <- de.update
    
    count <- count + 1
    
    #cat("Iteration ", count," done. Current tolerance is ",tol," and took ",proc.time()[3]-iter.start," seconds.","\n")
    #cat("The true value is ",results$x.current[1]," and the bootstrap value is ",x.current[1],"\n")
    
    
  }
  gc(verbose = FALSE)
  (c(x.current, de.current))
}

res_tmp <- lapply(1:100, boot_maker)

results_boot <- do.call(rbind, res_tmp)


save(results_boot, file=paste0(getwd(), "/EMestimates_boot.Rda"))

cat("Replication file 2 done.","\n")
